#include <stdio.h>
#include <stdlib.h>
#include <pthread.h>
#include <string.h>
#include <unistd.h>

// 任务结构
typedef struct {
    void (*function)(void *);
    void *argument;
} ThreadPoolTask;

// 线程池结构
typedef struct {
    pthread_t *threads;
    ThreadPoolTask *taskQueue;
    int threadCount;
    int queueCapacity;
    int queueSize;
    int queueFront;
    int queueRear;
    pthread_mutex_t mutex;
    pthread_cond_t condition;
    int shutdown;
} ThreadPool;

// 工作线程函数
static void *threadPoolWorker(void *pool) {
    ThreadPool *threadPool = (ThreadPool *)pool;
    ThreadPoolTask task;

    while (1) {
        pthread_mutex_lock(&threadPool->mutex);

        // 等待任务或关闭信号
        while (threadPool->queueSize == 0 && !threadPool->shutdown) {
            pthread_cond_wait(&threadPool->condition, &threadPool->mutex);
        }

        // 如果线程池正在关闭，退出线程
        if (threadPool->shutdown) {
            pthread_mutex_unlock(&threadPool->mutex);
            pthread_exit(NULL);
        }

        // 从队列中获取任务
        task.function = threadPool->taskQueue[threadPool->queueFront].function;
        task.argument = threadPool->taskQueue[threadPool->queueFront].argument;

        // 更新队列指针
        threadPool->queueFront = (threadPool->queueFront + 1) % threadPool->queueCapacity;
        threadPool->queueSize--;

        pthread_mutex_unlock(&threadPool->mutex);

        // 执行任务
        (*task.function)(task.argument);
    }

    return NULL;
}

// 创建线程池
ThreadPool *threadPoolCreate(int threadCount, int queueCapacity) {
    ThreadPool *pool;
    int i;

    // 参数检查
    if (threadCount <= 0 || queueCapacity <= 0) {
        return NULL;
    }

    // 分配线程池结构
    pool = (ThreadPool *)malloc(sizeof(ThreadPool));
    if (!pool) {
        return NULL;
    }

    // 初始化线程池结构
    pool->threadCount = threadCount;
    pool->queueCapacity = queueCapacity;
    pool->queueSize = 0;
    pool->queueFront = 0;
    pool->queueRear = 0;
    pool->shutdown = 0;

    // 分配线程数组和任务队列
    pool->threads = (pthread_t *)malloc(sizeof(pthread_t) * threadCount);
    pool->taskQueue = (ThreadPoolTask *)malloc(sizeof(ThreadPoolTask) * queueCapacity);

    if (!pool->threads || !pool->taskQueue) {
        free(pool->threads);
        free(pool->taskQueue);
        free(pool);
        return NULL;
    }

    // 初始化互斥锁和条件变量
    pthread_mutex_init(&pool->mutex, NULL);
    pthread_cond_init(&pool->condition, NULL);

    // 创建工作线程
    for (i = 0; i < threadCount; i++) {
        if (pthread_create(&pool->threads[i], NULL, threadPoolWorker, pool) != 0) {
            // 创建线程失败，清理资源
            threadPoolDestroy(pool);
            return NULL;
        }
    }

    return pool;
}

// 添加任务到线程池
int threadPoolAddTask(ThreadPool *pool, void (*function)(void *), void *argument) {
    if (!pool || !function) {
        return -1;
    }

    pthread_mutex_lock(&pool->mutex);

    // 检查线程池是否已满
    if (pool->queueSize >= pool->queueCapacity) {
        pthread_mutex_unlock(&pool->mutex);
        return -1;
    }

    // 检查线程池是否正在关闭
    if (pool->shutdown) {
        pthread_mutex_unlock(&pool->mutex);
        return -1;
    }

    // 添加任务到队列
    pool->taskQueue[pool->queueRear].function = function;
    pool->taskQueue[pool->queueRear].argument = argument;
    pool->queueRear = (pool->queueRear + 1) % pool->queueCapacity;
    pool->queueSize++;

    // 通知等待的线程
    pthread_cond_signal(&pool->condition);
    pthread_mutex_unlock(&pool->mutex);

    return 0;
}

// 获取线程池中的任务数量
int threadPoolGetTaskCount(ThreadPool *pool) {
    if (!pool) {
        return -1;
    }

    int count;
    pthread_mutex_lock(&pool->mutex);
    count = pool->queueSize;
    pthread_mutex_unlock(&pool->mutex);

    return count;
}

// 销毁线程池
void threadPoolDestroy(ThreadPool *pool) {
    if (!pool) {
        return;
    }

    pthread_mutex_lock(&pool->mutex);
    pool->shutdown = 1;
    pthread_mutex_unlock(&pool->mutex);

    // 唤醒所有等待的线程
    pthread_cond_broadcast(&pool->condition);

    // 等待所有线程结束
    for (int i = 0; i < pool->threadCount; i++) {
        pthread_join(pool->threads[i], NULL);
    }

    // 释放资源
    free(pool->threads);
    free(pool->taskQueue);
    pthread_mutex_destroy(&pool->mutex);
    pthread_cond_destroy(&pool->condition);
    free(pool);
}